﻿#include <stdio.h>
#include "Interpreter.hpp"
#include <iostream>
#include <opencv2/opencv.hpp>
#include <stdio.h>
#include "direct.h"
#include "io.h"
#include "fstream"

using namespace MNN;
using namespace cv;
using namespace std;

typedef struct DetectionRes {
	float x, y, w, h, prob;
} DetectionRes;

float sigmoid(float in) {
	return 1.f / (1.f + exp(-in));
}

void getFiles(string path, vector<string>& files)
{
	//文件句柄
	intptr_t   hFile = 0;
	//文件信息
	struct _finddata_t fileinfo;
	string p;
	if ((hFile = _findfirst(p.assign(path).append("\\*").c_str(), &fileinfo)) != -1)
	{
		do
		{
			//如果是目录,迭代之
			//如果不是,加入列表
			if ((fileinfo.attrib &  _A_SUBDIR))
			{
				if (strcmp(fileinfo.name, ".") != 0 && strcmp(fileinfo.name, "..") != 0)
					getFiles(p.assign(path).append("\\").append(fileinfo.name), files);
			}
			else
			{
				files.push_back(p.assign(path).append("\\").append(fileinfo.name));
			}
		} while (_findnext(hFile, &fileinfo) == 0);
		_findclose(hFile);
	}
}

// prepare img
cv::Mat prepareImage(cv::Mat& img, int input_c, int input_w, int input_h) 
{
	cv::Mat rgb;
	cv::cvtColor(img, rgb, COLOR_BGR2RGB);

	float scale = min(float(input_w) / img.cols, float(input_h) / img.rows);
	auto scaleSize = cv::Size(img.cols * scale, img.rows * scale);
	
	cv::Mat resized;
	cv::resize(rgb, resized, scaleSize, 0, 0, INTER_CUBIC);

	cv::Mat cropped(input_h, input_w, CV_8UC3, 127);
	Rect rect((input_w - scaleSize.width) / 2, (input_h - scaleSize.height) / 2, scaleSize.width, scaleSize.height);
	resized.copyTo(cropped(rect));
	cv::Mat img_float;
	cropped.convertTo(img_float, CV_32FC3, 1.f / 255.0);

	return img_float;
}

static int EntryIndex(int side, int lcoords, int lclasses, int location, int entry) {
	int n = location / (side * side);
	int loc = location % (side * side);
	return n * side * side * (lcoords + lclasses + 1) + entry * side * side + loc;
}

void DoNms(vector<DetectionRes>& detections, float nmsThresh) {
	auto iouCompute = [](float * lbox, float* rbox) {
		float interBox[] = {
			max(lbox[0], rbox[0]), //left
			min(lbox[0] + lbox[2], rbox[0] + rbox[2]), //right
			max(lbox[1], rbox[1]), //top
			min(lbox[1] + lbox[3], rbox[1] + rbox[3]), //bottom
		};

		if (interBox[2] >= interBox[3] || interBox[0] >= interBox[1])
			return 0.0f;

		float interBoxS = (interBox[1] - interBox[0] + 1) * (interBox[3] - interBox[2] + 1);
		return interBoxS / (lbox[2] * lbox[3] + rbox[2] * rbox[3] - interBoxS);
	};

	sort(detections.begin(), detections.end(), [=](const DetectionRes & left, const DetectionRes & right) {
		return left.prob > right.prob;
	});

	vector<DetectionRes> result;
	for (unsigned int m = 0; m < detections.size(); ++m) {
		result.push_back(detections[m]);
		for (unsigned int n = m + 1; n < detections.size(); ++n) {
			if (iouCompute((float *)(&detections[m]), (float *)(&detections[n])) > nmsThresh) {
				detections.erase(detections.begin() + n);
				--n;
			}
		}
	}
	detections = move(result);
}

int main()
{
	//string model_name;
	//string strImageFile;
	//string strImageFile_save;
	//cout << "please xxx.mnn file path" << endl;
	//cin >> model_name;
	//cout << "please img file path" << endl;
	//cin >> strImageFile;
	//cout << "please input save file path" << endl;
	//cin >> strImageFile_save;

	std::string model_name = "../model/yolov3-tiny.mnn";//模型地址
	std::string strImageFile = "../data";//图片地址
	string strImageFile_save = "../save_end_img\\";
	float FLAGS_t = 0.30;
	float FLAGS_iou_t = 0.45;
	int threads = 4;//计划使用并发数
	std::vector<float> anchors = { 10,14, 23,27, 37,58, 81,82, 135,169, 344,319 };

	std::shared_ptr<Interpreter> net(Interpreter::createFromFile(model_name.c_str()));//模型翻译创建
	ScheduleConfig config;//计划配置
	config.numThread = threads;
	config.type = MNN_FORWARD_CPU;
	auto session = net->createSession(config);
	auto input = net->getSessionInput(session, NULL);

	std::vector<int> shape = input->shape();
	int input_C = shape[1];
	int input_H = shape[2];
	int input_W = shape[3];
	//cout << input_C << "    " << input_H << "     " << input_W;
	vector<string> files;
	getFiles(strImageFile, files);
	char str[30];
	int size = files.size();
	for (int img_num = 0; img_num < size; img_num++)
	{
		cv::Mat image;
		cv::Mat img = cv::imread(files[img_num]);
		auto t_start_pre = std::chrono::high_resolution_clock::now();
		image = prepareImage(img,input_C,input_H,input_W);
		auto t_end_pre = std::chrono::high_resolution_clock::now();
		float total_pre = std::chrono::duration<float, std::milli>(t_end_pre - t_start_pre).count();
		std::cout << "prepare image take: " << total_pre << " ms." << endl;
		int r = image.rows;

		std::vector<int> dims{1, input_H, input_W, input_C};
		auto nchw_Tensor = MNN::Tensor::create<float>(dims, NULL, MNN::Tensor::TENSORFLOW);
		auto nchw_data = nchw_Tensor->host<float>();
		auto nchw_size = nchw_Tensor->size();
		::memcpy(nchw_data, image.data, nchw_size);
		input->copyFromHostTensor(nchw_Tensor);
		
		std::vector<int> shape = input->shape();
		int input_n = shape[0];
		int input_C = shape[1];
		int input_H = shape[2];
		int input_W = shape[3];
		cout <<input_n<<"        " <<input_C << "    " << input_H << "     " << input_W;
		
		// run network
		clock_t start = clock(); //计时
		net->runSession(session); //运行
		clock_t end = clock(); //计时
		double tim = end - start;
		cout << tim << "\n";

		std::string output_tensor13 = "016_convolutional";
		std::string output_tensor26 = "023_convolutional";

		auto tensor_scores13 = net->getSessionOutput(session, output_tensor13.c_str());
		auto tensor_scores26 = net->getSessionOutput(session, output_tensor26.c_str());

		auto dimType = tensor_scores13->getDimensionType();
		if (tensor_scores13->getType().code != halide_type_float) {
			dimType = Tensor::TENSORFLOW;
		}

		std::shared_ptr<Tensor> outputUser13(new Tensor(tensor_scores13, dimType));
		tensor_scores13->copyToHostTensor(outputUser13.get());

		std::shared_ptr<Tensor> outputUser26(new Tensor(tensor_scores26, dimType));
		tensor_scores26->copyToHostTensor(outputUser26.get());

		auto size13 = outputUser13->elementSize();
		auto size26 = outputUser26->elementSize();

		auto float_scores13 = outputUser13->host<float>();
		auto float_scores26 = outputUser26->host<float>();

 		std::vector<DetectionRes> objects;
		int out_c = tensor_scores13->channel();
		int classes = out_c/ 3 - 5;

		// postprocess
		int side13 = tensor_scores13->height();
		int side_square13 = side13 * side13;
		for (int i13 = 0; i13 < side_square13; ++i13) {
			int row13 = i13 / side13;
			int col13 = i13 % side13;
			for (int n13 = 0; n13 < 3; ++n13) {
				int obj_index13 = EntryIndex(side13, 4, classes, n13 * side13 * side13 + i13, 4);
				int box_index13 = EntryIndex(side13, 4, classes, n13 * side13 * side13 + i13, 0);
				float scale13 = float_scores13[obj_index13];
				if (scale13 < FLAGS_t)
					continue;
				double x13 = (col13 + sigmoid(float_scores13[box_index13 + 0 * side_square13])) / side13;
				double y13 = (row13 + sigmoid(float_scores13[box_index13 + 1 * side_square13])) / side13;
				double height13 = std::exp(float_scores13[box_index13 + 3 * side_square13]) * anchors[6 + 2 * n13 + 1]/input_H;
				double width13 = std::exp(float_scores13[box_index13 + 2 * side_square13]) * anchors[6 + 2 * n13]/input_W;
				double xn13 = x13 - width13 / 2;
				double yn13 = y13 - height13 / 2;
				int label13 = 0;
				int name13 = 1;
				float score13 = 0;
				for (int j13 = 0; j13 < classes; ++j13)
				{
					int class_index13 = EntryIndex(side13, 4, classes, n13 * side_square13 + i13, 4 + 1 + j13);
					float prob13 = sigmoid(scale13) * sigmoid(float_scores13[class_index13]);
					if (prob13 > score13)
					{
						score13 = prob13;
						label13 = name13;
						name13 += 1;
					}
					else
					{
						name13 += 1;
					}
				}
				if (score13 < FLAGS_t)
					continue;
				DetectionRes det;
				det.x = xn13;
				det.y = yn13;
				det.w = width13;
				det.h = height13;
				det.prob = label13;
				objects.push_back(det);
			}
		}
		int side26 = tensor_scores26->height();
		int side_square26 = side26 * side26;
		for (int i26 = 0; i26 < side_square26; ++i26) {
			int row26 = i26 / side26;
			int col26 = i26 % side26;
			for (int n26 = 0; n26 < 3; ++n26) {
				int obj_index26 = EntryIndex(side26, 4, classes, n26 * side26 * side26 + i26, 4);
				int box_index26 = EntryIndex(side26, 4, classes, n26 * side26 * side26 + i26, 0);
				float scale26 = float_scores26[obj_index26];
				if (scale26 < FLAGS_t)
					continue;
				double x26 = (col26 + sigmoid(float_scores26[box_index26 + 0 * side_square26])) / side26;
				double y26 = (row26 + sigmoid(float_scores26[box_index26 + 1 * side_square26])) / side26;
				double height26 = std::exp(float_scores26[box_index26 + 3 * side_square26]) * anchors[2 * n26 + 1]/input_H;
				double width26 = std::exp(float_scores26[box_index26 + 2 * side_square26]) * anchors[2 * n26]/input_W;
				//cout << x << "   " << y << "  " << height << "  " << width<<endl;
				double xn26 = x26 - width26 / 2;
				double yn26 = y26 - height26 / 2;
				int label26 = 0;
				int name26 = 1;
				float score26 = 0;
				for (int j26 = 0; j26 < classes; ++j26)
				{
					int class_index26 = EntryIndex(side26, 4, classes, n26 * side_square26 + i26, 4 + 1 + j26);
					float prob26 = sigmoid(scale26) * sigmoid(float_scores26[class_index26]);
					class_index26 = EntryIndex(side26, 4, classes, n26 * side_square26 + i26, 4 + 1 + j26);
					prob26 = sigmoid(scale26) * sigmoid(float_scores26[class_index26]);
					if (prob26 > score26)
					{
						score26 = prob26;
						label26 = name26;
						name26 += 1;
					}
					else
					{
						name26 += 1;
					}
				}
				if (score26 < FLAGS_t)
					continue;
				DetectionRes det;
				det.x = xn26;
				det.y = yn26;
				det.w = width26;
				det.h = height26;
				det.prob = label26;
				objects.push_back(det);
			}
		}
		int width = img.cols;
		int height = img.rows;
		float scale = min(float(input_W) / width, float(input_H) / height);
		float scaleSize[] = { width * scale, height * scale };

		//correct box
		for (auto& bbox : objects) {
			bbox.x = (bbox.x * input_W - (input_W - scaleSize[0]) / 2.f) / scale;
			bbox.y = (bbox.y * input_H - (input_H - scaleSize[1]) / 2.f) / scale;
			bbox.w *= input_W;
			bbox.h *= input_H;
			bbox.w /= scale;
			bbox.h /= scale;
		}

		float nmsThresh = FLAGS_iou_t;
		if (nmsThresh > 0)
			DoNms(objects, nmsThresh);

		for (int i = 0; i < objects.size(); ++i)
		{
			int x = objects[i].x,
				y = objects[i].y,
				w = objects[i].w,
				h = objects[i].h,
			    label = objects[i].prob;
			//cout << x << "," << y << "," << w << "," << h <<","<<label<<","<< endl;
			cv::Rect rect = { x, y, w, h };
			cv::rectangle(img, rect, cv::Scalar(255, 0, 0), 2);
			cv::putText(img, std::to_string(label), Point(x, y), cv::FONT_HERSHEY_COMPLEX, 1, (255, 0, 0), 1);
			//cv::putText(img, std::to_string(conf), Point(x, y + 30), cv::FONT_HERSHEY_COMPLEX, 1, (255, 0, 0), 1);
		}
		//imshow("matShowResult",src);
		//cv::waitKey(0);
		cv::imwrite(strImageFile_save + std::to_string(img_num) + ".jpg", img);
	}
	system("pause");
}
